package com.ruoyi.pac4j;

import cc.renzhihao.sso.bean.JwtAuthenticator;
import cc.renzhihao.sso.bean.JwtClient;
import cc.renzhihao.sso.enums.AuthWayEnum;
import cc.renzhihao.sso.pac4j.GenericOAuth20Client;
import cc.renzhihao.sso.pac4j.IAMSAML2IdentityProviderMetadataResolver;
import cn.hutool.core.net.url.UrlBuilder;
import cn.hutool.json.JSONUtil;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.github.scribejava.core.model.Verb;
import com.google.common.collect.Lists;
import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.util.Base64;
import com.nimbusds.oauth2.sdk.auth.ClientAuthenticationMethod;
import com.ruoyi.common.utils.file.FileUploadUtils;
import com.ruoyi.framework.config.RuoYiConfig;
import com.ruoyi.project.system.auth.domain.SysAuth;
import com.ruoyi.project.system.auth.service.AuthService;
import com.ruoyi.project.system.keystore.domain.Keystore;
import com.ruoyi.project.system.keystore.service.KeystoreService;
import org.apache.commons.lang3.StringUtils;
import org.opensaml.saml.saml2.core.Assertion;
import org.opensaml.saml.saml2.core.Attribute;
import org.opensaml.saml.saml2.core.AuthnStatement;
import org.opensaml.saml.saml2.core.Response;
import org.pac4j.cas.client.CasClient;
import org.pac4j.cas.config.CasConfiguration;
import org.pac4j.core.client.Client;
import org.pac4j.core.client.Clients;
import org.pac4j.core.config.Config;
import org.pac4j.core.profile.ProfileHelper;
import org.pac4j.jwt.config.encryption.SecretEncryptionConfiguration;
import org.pac4j.jwt.config.signature.SecretSignatureConfiguration;
import org.pac4j.oauth.client.CasOAuthWrapperClient;
import org.pac4j.oauth.profile.JsonHelper;
import org.pac4j.oauth.profile.casoauthwrapper.CasOAuthWrapperProfile;
import org.pac4j.oauth.profile.casoauthwrapper.CasOAuthWrapperProfileDefinition;
import org.pac4j.oidc.client.OidcClient;
import org.pac4j.oidc.config.OidcConfiguration;
import org.pac4j.oidc.redirect.OidcRedirectionActionBuilder;
import org.pac4j.saml.client.SAML2Client;
import org.pac4j.saml.config.SAML2Configuration;
import org.pac4j.saml.context.SAML2MessageContext;
import org.pac4j.saml.credentials.SAML2Credentials;
import org.pac4j.saml.profile.converter.SimpleSAML2AttributeConverter;
import org.pac4j.saml.sso.impl.SAML2AuthnResponseValidator;
import org.pac4j.scribe.builder.api.CasOAuthWrapperApi20;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.DependsOn;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.FileSystemResource;
import org.springframework.stereotype.Component;

import java.io.File;
import java.util.*;

import static org.pac4j.core.profile.AttributeLocation.PROFILE_ATTRIBUTE;

@Component
@DependsOn("ruoYiConfig")
public class Pac4jService {

    private static final Logger log = LoggerFactory.getLogger(Pac4jService.class);

    @Autowired
    private AuthService authService;

    @Autowired
    private KeystoreService keystoreService;

    @Autowired
    private Config config;


    public void init() {
        SysAuth search = new SysAuth();
        search.setStatus(0);
        List<SysAuth> sysAuths = authService.selectAuthList(search);
        List<Client> clientList = Lists.newArrayList();

        for (SysAuth sysAuth : sysAuths) {
            if (StringUtils.equalsIgnoreCase(sysAuth.getType(), AuthWayEnum.CAS.getName())) {
                clientList.add(buildCasClient(sysAuth));
                continue;
            }

            if (StringUtils.equalsIgnoreCase(sysAuth.getType(), AuthWayEnum.OAUTH2.getName())) {
                clientList.add(buildOauthClient(sysAuth));
                continue;
            }

            if (StringUtils.equalsIgnoreCase(sysAuth.getType(), AuthWayEnum.OIDC.getName())) {
                clientList.add(buildOidcClient(sysAuth));
                continue;
            }

            if (StringUtils.equalsIgnoreCase(sysAuth.getType(), AuthWayEnum.JWT.getName())) {
                clientList.add(buildJwtClient(sysAuth));
                continue;
            }

            if (StringUtils.equalsIgnoreCase(sysAuth.getType(), "GENERIC_OAUTH2")) {
                clientList.add(buildGenericOauthClient(sysAuth));
                continue;
            }

            if (StringUtils.equalsIgnoreCase(sysAuth.getType(), AuthWayEnum.SAML.getName())) {
                clientList.add(buildSamlClient(sysAuth));
            }
        }

        Clients clients = new Clients(clientList);
        config.setCallbackLogic(new MyCallbackLogic(authService));
        config.setClients(clients);
    }

    private Client buildCasClient(SysAuth detail) {
        log.debug("添加casClient: {}", JSONUtil.toJsonStr(detail));
        CasClient casClient = new CasClient();
        CasConfiguration casConfiguration = new CasConfiguration();
        casConfiguration.setPrefixUrl(detail.getAuthUrl());
        casClient.setConfiguration(casConfiguration);
        casClient.setName(detail.getClientName());

        String callbackUrl = UrlBuilder.ofHttp(detail.getLocalServer()).addPath(detail.getCallbackUrl()).toString();
        casClient.setCallbackUrl(callbackUrl);
        return casClient;
    }


    private Client buildOauthClient(SysAuth detail) {
        log.debug("添加oauthClient: {}", JSONUtil.toJsonStr(detail));

        final CasOAuthWrapperClient oauthClient = new CasOAuthWrapperClient() {
            @Override
            protected void internalInit(boolean forceReinit) {
                super.internalInit(forceReinit);
                configuration.setProfileDefinition(new CasOAuthWrapperProfileDefinition() {
                    @Override
                    public CasOAuthWrapperProfile extractUserProfile(String body) {
                        final var profile = (CasOAuthWrapperProfile) newProfile();
                        var json = JsonHelper.getFirstNode(body);
                        if (json != null) {
                            profile.setId(ProfileHelper.sanitizeIdentifier(JsonHelper.getElement(json, "id")));
                            // CAS <= v4.2
                            if (json instanceof ArrayNode) {
                                for (JsonNode jsonNode : json) {
                                    json = jsonNode;
                                    final var attribute = json.fieldNames().next();
                                    convertAndAdd(profile, PROFILE_ATTRIBUTE, attribute, JsonHelper.getElement(json, attribute));
                                }
                                // CAS v5
                            } else if (json instanceof ObjectNode) {
                                final var keys = json.fieldNames();
                                while (keys.hasNext()) {
                                    final var key = keys.next();
                                    convertAndAdd(profile, PROFILE_ATTRIBUTE, key, JsonHelper.getElement(json, key));
                                }
                            }
                        } else {
                            raiseProfileExtractionJsonError(body);
                        }
                        return profile;
                    }
                });
                configuration.setApi(new CasOAuthWrapperApi20(getCasOAuthUrl(), true, Verb.POST));
            }
        };
        oauthClient.setName(detail.getClientName());
        oauthClient.setCasOAuthUrl(detail.getAuthUrl());
        oauthClient.setKey(detail.getClientId());
        oauthClient.setSecret(detail.getClientSecret());

        oauthClient.setAuthorizationGenerator((ctx, session, profile) -> Optional.of(profile));

        String callbackUrl = UrlBuilder.ofHttp(detail.getLocalServer()).addPath(detail.getCallbackUrl()).toString();
        oauthClient.setCallbackUrl(callbackUrl);
        return oauthClient;
    }

    private Client buildOidcClient(SysAuth detail) {
        log.debug("添加oidcClient: {}", JSONUtil.toJsonStr(detail));

        final OidcConfiguration oidcConfiguration = new OidcConfiguration();
        oidcConfiguration.setClientId(detail.getClientId());
        oidcConfiguration.setSecret(detail.getClientSecret());
        oidcConfiguration.setUseNonce(false);
        oidcConfiguration.setDiscoveryURI(detail.getWellKnown());
        oidcConfiguration.setScope(detail.getScope());
        oidcConfiguration.setClientAuthenticationMethod(ClientAuthenticationMethod.CLIENT_SECRET_BASIC);
        oidcConfiguration.addCustomParam("prompt", "consent");
        oidcConfiguration.setWithState(false);
        oidcConfiguration.setPreferredJwsAlgorithm(JWSAlgorithm.RS256);

        final OidcClient oidcClient = new OidcClient(oidcConfiguration);
        oidcClient.setName(detail.getClientName());
        oidcClient.setRedirectionActionBuilder(new OidcRedirectionActionBuilder(oidcClient));

        oidcClient.setAuthorizationGenerator((ctx, session, profile) -> Optional.of(profile));

        String callbackUrl = UrlBuilder.ofHttp(detail.getLocalServer()).addPath(detail.getCallbackUrl()).toString();
        oidcClient.setCallbackUrl(callbackUrl);
        return oidcClient;
    }


    public Client buildGenericOauthClient(SysAuth detail) {
        log.debug("添加genericOauthClient: {}", JSONUtil.toJsonStr(detail));
        String authUrl = detail.getAuthUrl();
        String tokenUrl = detail.getTokenUrl();
        String tokenParamPosition = detail.getTokenParamPosition();
        final GenericOAuth20Client oauthClient = new GenericOAuth20Client();
        oauthClient.setName(detail.getClientName());
        oauthClient.setKey(detail.getClientId());
        oauthClient.setSecret(detail.getClientSecret());
        oauthClient.setAuthUrl(authUrl);
        oauthClient.setTokenUrl(tokenUrl);
        oauthClient.setTokenParamPosition(tokenParamPosition);
        oauthClient.setProfileNodePath(detail.getProfilePath());
        oauthClient.setProfileId(detail.getProfileId());
        oauthClient.setProfileUrl(detail.getProfileUrl());
        Verb profileVerb;
        try {
            profileVerb = Verb.valueOf(Optional.ofNullable(detail.getProfileVerb()).orElse("GET").toUpperCase(Locale.ROOT));
        } catch (IllegalArgumentException e) {
            profileVerb = Verb.GET;
        }
        oauthClient.setProfileVerb(profileVerb);
        oauthClient.setScope(detail.getScope());

        oauthClient.setAuthorizationGenerator((ctx, session, profile) -> Optional.of(profile));

        String callbackUrl = UrlBuilder.ofHttp(detail.getLocalServer()).addPath(detail.getCallbackUrl()).toString();
        oauthClient.setCallbackUrl(callbackUrl);
        return oauthClient;
    }

    private Client buildJwtClient(SysAuth detail) {
        log.debug("添加jwtClient: {}", JSONUtil.toJsonStr(detail));

        final CasConfiguration configuration = new CasConfiguration();
        configuration.setPrefixUrl(detail.getAuthUrl());
        final JwtAuthenticator authenticator = new JwtAuthenticator(new SecretSignatureConfiguration(detail.getSignatureKey()),
                new SecretEncryptionConfiguration(new Base64(detail.getEncryptionKey()).decode()));
        final JwtClient jwtClient = new JwtClient(configuration, authenticator);
        jwtClient.setName(detail.getClientName());

        String callbackUrl = UrlBuilder.ofHttp(detail.getLocalServer()).addPath(detail.getCallbackUrl()).toString();
        jwtClient.setCallbackUrl(callbackUrl);
        return jwtClient;
    }

    private Client buildSamlClient(SysAuth detail) {
        log.debug("添加samlClient: {}", JSONUtil.toJsonStr(detail));

        Long keystoreId = detail.getKeystoreId() == null ? -1L : detail.getKeystoreId();
        Keystore keystore = keystoreService.selectKeystoreById(keystoreId);


        final SAML2Configuration configuration = new SAML2Configuration();

        if (keystore == null || Objects.equals(keystore.getInitial(), 0)) {
            configuration.setKeystoreResource(new ClassPathResource("samlKeystore.jks"));
            configuration.setKeystorePassword("pac4j-demo-passwd");
            configuration.setPrivateKeyPassword("pac4j-demo-passwd");
        } else {
            String fileName = keystore.getFileName();
            String filePath = RuoYiConfig.getUploadPath();
            String absPath = FileUploadUtils.getAbsoluteFile(filePath, fileName).getAbsolutePath();

            configuration.setKeystoreResource(new FileSystemResource(absPath));
            configuration.setKeystorePassword(keystore.getKeystorePass());
            configuration.setPrivateKeyPassword(keystore.getPrivateKeyPass());
        }

        configuration.setIdentityProviderMetadataPath(detail.getMetadataPath());
        configuration.setServiceProviderEntityId(detail.getEntityId());
        configuration.setMaximumAuthenticationLifetime(3600);

        String localPath = RuoYiConfig.getProfile() + "/saml/";
        configuration.setServiceProviderMetadataResource(new FileSystemResource(new File(localPath + detail.getMetadataFilename()).getAbsoluteFile()));
        configuration.setForceServiceProviderMetadataGeneration(true);
        configuration.setNameIdAttribute("account");
        configuration.setAllSignatureValidationDisabled(StringUtils.equals(detail.getAllSignatureValidationDisabled(), "Y"));
        configuration.setAuthnRequestSigned(StringUtils.equals(detail.getAuthnRequestSigned(), "Y"));
        configuration.setSpLogoutRequestSigned(StringUtils.equals(detail.getSpLogoutRequestSigned(), "Y"));
        configuration.setWantsAssertionsSigned(StringUtils.equals(detail.getWantsAssertionsSigned(), "Y"));
        configuration.setWantsResponsesSigned(StringUtils.equals(detail.getWantsResponsesSigned(), "Y"));
        configuration.setIdentityProviderMetadataResolver(new IAMSAML2IdentityProviderMetadataResolver(configuration));
        final SAML2Client saml2Client = new SAML2Client(configuration) {
            @Override
            protected void initSAMLResponseValidator() {
                // Build the SAML response validator
                this.authnResponseValidator = new SAML2AuthnResponseValidator(
                        this.signatureTrustEngineProvider,
                        this.decrypter,
                        this.replayCache,
                        this.configuration) {
                    @Override
                    protected SAML2Credentials buildSAML2Credentials(SAML2MessageContext context, Response response) {
                        final Assertion subjectAssertion = context.getSubjectAssertion();

                        final List<Attribute> attributes = collectAssertionAttributes(subjectAssertion);
                        final SAML2Credentials.SAMLNameID samlNameId = determineNameID(context, from(attributes));

                        final String sessionIndex = getSessionIndex(subjectAssertion);
                        final String sloKey = computeSloKey(sessionIndex, samlNameId);
                        if (sloKey != null) {
                            logoutHandler.recordSession(context.getWebContext(), context.getSessionStore(), sloKey);
                        }

                        final String issuerEntityId = subjectAssertion.getIssuer().getValue();
                        final List<AuthnStatement> authnStatements = subjectAssertion.getAuthnStatements();
                        final List<String> authnContexts = new ArrayList<>();
                        for (final AuthnStatement authnStatement : authnStatements) {
                            if (authnStatement.getAuthnContext().getAuthnContextClassRef() != null) {
                                authnContexts.add(authnStatement.getAuthnContext().getAuthnContextClassRef().getURI());
                            }
                        }
                        return new SAML2Credentials(samlNameId, issuerEntityId, from(attributes),
                                subjectAssertion.getConditions(), sessionIndex, authnContexts, "");
                    }

                    public List<SAML2Credentials.SAMLAttribute> from(final List<Attribute> samlAttributes) {
                        return SAML2Credentials.SAMLAttribute.from(new SimpleSAML2AttributeConverter(), samlAttributes);
                    }
                };
                this.authnResponseValidator.setAcceptedSkew(this.configuration.getAcceptedSkew());
            }
        };
        saml2Client.setName(detail.getClientName());

        String callbackUrl = UrlBuilder.ofHttp(detail.getLocalServer()).addPath(detail.getCallbackUrl()).toString();
        saml2Client.setCallbackUrl(callbackUrl);
        return saml2Client;
    }

}
